import cv2
import os
import numpy as np


def extract_seal(image_path, output_path=None, h_scale=[160,10],s_scale=[50,255],v_scale=[50,255]):
    """
    提取图像中的印章区域
    :param image_path: 输入图像路径
    :param output_path: 输出图像路径（None则不保存）
    :param h_scale: 色调范围 [min, max]
    :param s_scale: 饱和度范围 [min, max]
    :param v_scale: 明度范围 [min, max]
    :return:
    """

    hsv_lower = [h_scale[0], s_scale[0], v_scale[0]]
    hsv_upper = [h_scale[1], s_scale[1], v_scale[1]]
    img = extract_red_seal(image_path, output_path, hsv_lower, hsv_upper)
    return img
def extract_red_seal(image_path, output_path=None, hsv_lower=[0, 50, 50], hsv_upper=[10, 255, 255]):
    """
    提取图像中的红色印章区域
    :param image_path: 输入图像路径
    :param output_path: 输出图像路径（None则不保存）
    :param red_threshold: 红色通道阈值（0-255）
    :param hsv_lower: HSV颜色空间下限 [H_min, S_min, V_min]
    :param hsv_upper: HSV颜色空间上限 [H_max, S_max, V_max]
    :return:
    """

    image = cv2.imread(image_path)
    hue_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    if(hsv_lower[0]<hsv_upper[0]):
        lower_red1 = np.array(hsv_lower)
        upper_red1 = np.array(hsv_upper)
        red_mask = cv2.inRange(hue_image, lower_red1, upper_red1)
    else:
        lower_red1 = np.array([0, hsv_lower[1], hsv_lower[2]])
        upper_red1 = np.array(hsv_upper)
        lower_red2 = np.array(hsv_lower)
        upper_red2 = np.array([180, hsv_upper[1], hsv_upper[2]]);
        mask1 = cv2.inRange(hue_image, lower_red1, upper_red1)
        mask2 = cv2.inRange(hue_image, lower_red2, upper_red2)
        red_mask = cv2.bitwise_or(mask1, mask2)

    # lower_red1 = np.array([0, 20, 20])  # 降低S和V的最小值
    # upper_red1 = np.array([10, 255, 255])
    # lower_red2 = np.array([160, 20, 20])  # 同上
    # upper_red2 = np.array([180, 255, 255])
    # mask1 = cv2.inRange(hue_image, lower_red1, upper_red1)
    # mask2 = cv2.inRange(hue_image, lower_red2, upper_red2)
    # red_mask=cv2.bitwise_or(mask1,mask2)

    # 2. 检测黑色区域
    lower_black = np.array([0, 0, 0])
    upper_black = np.array([180, 255, 50])
    black_mask = cv2.inRange(hue_image, lower_black, upper_black)

    # red_mask = cv2.bitwise_and(red_mask, cv2.bitwise_not(black_mask))

    # 3. 形态学处理
    # 闭运算（先膨胀后腐蚀，填充小孔）
    # kernel = np.ones((3, 3), np.uint8)
    # red_mask = cv2.morphologyEx(red_mask, cv2.MORPH_CLOSE, kernel)
    # # 开运算（先腐蚀后膨胀，去除小噪点）
    # red_mask = cv2.morphologyEx(red_mask, cv2.MORPH_OPEN, kernel)

    # 4. 保存红色掩码为独立文件（格式：原图名称-redmask.png）
    if output_path:
        base_name = os.path.splitext(image_path)[0]  # 获取无扩展名的文件名
        redmask_path = os.path.join(output_path, f"{base_name}-redmask.png")
        cv2.imwrite(redmask_path, red_mask)
        print(f"Red mask saved to: {redmask_path}")

    index1 = red_mask == 255
    img = np.zeros(image.shape, np.uint8)
    img[:, :] = (255, 255, 255)
    img[index1] = image[index1]  # (0,0,255)

    if output_path:
        cv2.imwrite(output_path, img)
        print(f"Red seal extracted to: {output_path}")

    return img

def extract_seal_array(image,image_path=None,output_path=None, h_scale=[160,10],s_scale=[50,255],v_scale=[50,255]):
    hsv_lower = [h_scale[0], s_scale[0], v_scale[0]]
    hsv_upper = [h_scale[1], s_scale[1], v_scale[1]]
    hue_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    if (hsv_lower[0] < hsv_upper[0]):
        lower_red1 = np.array(hsv_lower)
        upper_red1 = np.array(hsv_upper)
        red_mask = cv2.inRange(hue_image, lower_red1, upper_red1)
    else:
        lower_red1 = np.array([0, hsv_lower[1], hsv_lower[2]])
        upper_red1 = np.array(hsv_upper)
        lower_red2 = np.array(hsv_lower)
        upper_red2 = np.array([180, hsv_upper[1], hsv_upper[2]]);
        mask1 = cv2.inRange(hue_image, lower_red1, upper_red1)
        mask2 = cv2.inRange(hue_image, lower_red2, upper_red2)
        red_mask = cv2.bitwise_or(mask1, mask2)


    # 2. 检测黑色区域
    lower_black = np.array([0, 0, 0])
    upper_black = np.array([180, 255, 50])
    black_mask = cv2.inRange(hue_image, lower_black, upper_black)

    # red_mask = cv2.bitwise_and(red_mask, cv2.bitwise_not(black_mask))

    # 3. 形态学处理
    # 闭运算（先膨胀后腐蚀，填充小孔）
    # kernel = np.ones((3, 3), np.uint8)
    # red_mask = cv2.morphologyEx(red_mask, cv2.MORPH_CLOSE, kernel)
    # # 开运算（先腐蚀后膨胀，去除小噪点）
    # red_mask = cv2.morphologyEx(red_mask, cv2.MORPH_OPEN, kernel)

    # 4. 保存红色掩码为独立文件（格式：原图名称-redmask.png）
    if output_path and image_path :
        base_name = os.path.splitext(image_path)[0]  # 获取无扩展名的文件名
        redmask_path = os.path.join(output_path, f"{base_name}-redmask.png")
        cv2.imwrite(redmask_path, red_mask)
        print(f"Red mask saved to: {redmask_path}")

    index1 = red_mask == 255
    img = np.zeros(image.shape, np.uint8)
    img[:, :] = (255, 255, 255)
    img[index1] = image[index1]  # (0,0,255)

    if output_path and image_path:
        cv2.imwrite(output_path, img)
        print(f"Red seal extracted to: {output_path}")

    return img




def batch_extract_seals(input_dir, output_dir,h_scale=[160,10],s_scale=[50,255],v_scale=[50,255]):

    '''
    :param input_dir:
    :param output_dir:
    :param h_scale:
    :param s_scale:
    :param v_scale:
    :return:
    '''
    # 创建输出文件夹
    os.makedirs(output_dir, exist_ok=True)

    # 获取所有图片文件
    valid_exts = ('.jpg', '.jpeg', '.png', '.bmp')
    image_files = [f for f in os.listdir(input_dir)
                   if f.lower().endswith(valid_exts)]

    for img_file in image_files:
        # 输入输出路径
        input_path = os.path.join(input_dir, img_file)
        base_name = os.path.splitext(img_file)[0]
        output_path = os.path.join(output_dir, f"{base_name}-extract.jpg")

        try:
            extract_seal(input_path, output_path, h_scale,s_scale,v_scale)

        except Exception as e:
            print(f"❌ 处理失败 {img_file}: {str(e)}")

def batch_extract_red_seals(input_dir, output_dir,
                            hsv_lower=[0, 50, 50], hsv_upper=[10, 255, 255]):
    """
    批量提取文件夹中所有图片的红色印章
    :param input_dir: 输入图片文件夹路径
    :param output_dir: 输出文件夹路径
    :param red_threshold: 红色通道阈值
    :param hsv_lower: HSV下限 [H,S,V]
    :param hsv_upper: HSV上限 [H,S,V]
    """
    # 创建输出文件夹
    os.makedirs(output_dir, exist_ok=True)

    # 获取所有图片文件
    valid_exts = ('.jpg', '.jpeg', '.png', '.bmp')
    image_files = [f for f in os.listdir(input_dir)
                   if f.lower().endswith(valid_exts)]

    for img_file in image_files:
        # 输入输出路径
        input_path = os.path.join(input_dir, img_file)
        base_name = os.path.splitext(img_file)[0]
        output_path = os.path.join(output_dir, f"{base_name}-extract.jpg")

        try:
            extract_red_seal(input_path, output_path,  hsv_lower, hsv_upper)

        except Exception as e:
            print(f"❌ 处理失败 {img_file}: {str(e)}")


# 使用示例
if __name__ == "__main__":
    # input_image = "seal_examples/example1.jpg"
    input_image = "cropped_seals"
    output_image = "cropped_seals_extract"



    mask = batch_extract_red_seals(
        input_dir=input_image,
        output_dir=output_image,
        hsv_lower=[160, 20, 20],  # 更严格的HSV范围
        hsv_upper=[10, 255, 255]
    )

    # 可视化结果（可选）
    # cv2.imshow("Red Seal Mask", mask)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
